import pymysql
from dbutils.pooled_db import PooledDB
from config import MYSQL_CONFIG


class MySQLPool:
    def __init__(self, host, user, password, database, port=3306, max_connections=10):
        self.pool = PooledDB(
            pymysql,
            maxconnections=max_connections,  # 最大连接数
            mincached=2,  # 初始化时，连接池中至少创建的空闲的连接
            maxcached=5,  # 最大空闲连接数
            maxshared=3,  # 最大共享连接数
            blocking=True,  # 连接池中如果没有可用连接后，是否阻塞
            setsession=[],
            host=host,
            user=user,
            password=password,
            database=database,
            port=port
        )

    def execute_query(self, query, params=None):
        """执行查询并返回结果"""
        connection = self.pool.connection()
        try:
            cursor = connection.cursor()
            cursor.execute(query, params)
            # 使用 yield 逐行返回结果
            # for data in cursor.fetchall():
            #     yield data  # 逐行生成数据
            results = cursor.fetchall()
            return results
        except Exception as e:
            print(f"Error executing query: {e}")
            return None
        finally:
            cursor.close()
            connection.close()  # 将连接归还给连接池

    def execute_update(self, query, params=None):
        """执行更新操作（INSERT/UPDATE/DELETE）"""
        connection = self.pool.connection()
        try:
            cursor = connection.cursor()
            cursor.execute(query, params)
            connection.commit()  # 提交事务
            return cursor.rowcount
        except Exception as e:
            print(f"Error executing update: {e}")
            return None
        finally:
            cursor.close()
            connection.close()  # 将连接归还给连接池


# 使用示例
if __name__ == "__main__":
    # 创建连接池
    db_pool = MySQLPool(host=MYSQL_CONFIG.get("host"), user=MYSQL_CONFIG.get("username"),
                        password=MYSQL_CONFIG.get("password"), database=MYSQL_CONFIG.get("database"),
                        port=MYSQL_CONFIG.get("port"),max_connections=5)

    # 查询示例
    select_query = "select id,source_link,download_status,download_fail_msg,transfer_status,transfer_error from " \
                   "storage_download WHERE user_id = %s and delete_time = 0 "
    result = db_pool.execute_query(select_query, '100304')
    dict_data = {}
    if result:
        for data in result:
            dict_data[data[1]] = {
                "download_id": data[0],
                "download_status": data[2],
                "transfer_status": data[4],
                "download_fail_msg": data[3],
                "transfer_error": data[5],
            }
    else:
        print("没有结果")
    print(dict_data.get("https://6tm.wdubo.com/20241105/H7GpUh6n/hls/index.m3u8?sign=Aia%2FC3eC%2BHYNwjBOLjZQmDK%2BUJUzs5LSayaOOBvo%2FfE%3D").get("download_id"))

    # 更新示例
    # update_query = "UPDATE your_table SET some_column = %s WHERE id = %s"
    # affected_rows = db_pool.execute_update(update_query, (new_value, record_id))
    # print(f"Affected rows: {affected_rows}")





